#!/usr/bin/env python36
# -*- coding: utf-8 -*-
"""
Created on July, 2018

@author: Tangrizzly
"""

import argparse
import pickle
import time
from utils import build_graph, Data, split_validation
from model import *

'''
参数申明：
    dataset='sample': 此处只用一个小的数据集用于检测代码的正确性
    batchSize=100: 批处理size设置100
    epoch=30: 在基于会话的推荐方案中大多数会话的长度相对较短，减少训练次数防止过拟合 
    hiddenSize=100: 基于其他论文设置的隐向量维度
    lr=0.001, lr_dc=0.1, lr_dc_step=31: 初始学习率设置为0.001，并且每3个周期将衰减0.1 
    ------Adam优化：结合RMSProp（结合梯度平方的指数移动平均数来调节学习率的变化）与AdaGrad自适应优化（对每个不同的参数调整不同的学习率，对频繁变化的参数以更小的步长进行更新，而稀疏的参数以更大的步长进行更新）。
    l2=1e-05: 权重衰减, L2正则化的目的就是为了让权重衰减到更小的值，在一定程度上减少模型过拟合的问题，所以权重衰减（weight decay）也叫L2正则化。
    nonhybrid=False: 只使用全局特征预测=false
    patience=10: the number of epoch to wait before early stop 提前停止之前的训练数？
    valid_portion=0.1: 将训练集的部分（10%）拆分为验证集
    validation=False: 批处理size设置100
    step=1: gnn传播步骤设置
'''
'''
argparse模块可以让人轻松编写用户友好的命令行接口。使用步骤：
    创建一个ArgumentParser对象
    调用 add_argument()方法添加此程序参数信息
    通过 parse_args() 方法解析参数
'''

# 设置参数
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='sample', help='dataset name: diginetica/yoochoose1_4/yoochoose1_64/sample')
parser.add_argument('--batchSize', type=int, default=100, help='input batch size')
parser.add_argument('--hiddenSize', type=int, default=100, help='hidden state size')
parser.add_argument('--epoch', type=int, default=30, help='the number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.001, help='learning rate')  # [0.001, 0.0005, 0.0001]
parser.add_argument('--lr_dc', type=float, default=0.1, help='learning rate decay rate')
parser.add_argument('--lr_dc_step', type=int, default=3, help='the number of steps after which the learning rate decay')
parser.add_argument('--l2', type=float, default=1e-5, help='l2 penalty')  # [0.001, 0.0005, 0.0001, 0.00005, 0.00001]
parser.add_argument('--step', type=int, default=1, help='gnn propogation steps')
parser.add_argument('--patience', type=int, default=10, help='the number of epoch to wait before early stop ')
# action='store_true'就代表着一旦有这个参数，做出动作“将其值标为True”
parser.add_argument('--nonhybrid', action='store_true', help='only use the global preference to predict')
parser.add_argument('--validation', action='store_true', help='validation')
parser.add_argument('--valid_portion', type=float, default=0.1,
                    help='split the portion of training set as validation set')
opt = parser.parse_args()
print(opt)


def main():
    # 加载数据
    train_data = pickle.load(open('../datasets/' + opt.dataset + '/train.txt', 'rb'))
    if opt.validation:
        train_data, valid_data = split_validation(train_data, opt.valid_portion)
        test_data = valid_data
    else:
        test_data = pickle.load(open('../datasets/' + opt.dataset + '/test.txt', 'rb'))
    # all_train_seq = pickle.load(open('../datasets/' + opt.dataset + '/all_train_seq.txt', 'rb'))
    # g = build_graph(all_train_seq)
    # 加载数据
    train_data = Data(train_data, shuffle=True)
    test_data = Data(test_data, shuffle=False)
    # del all_train_seq, g
    # 这里节点个数固定死，不太好啊
    if opt.dataset == 'diginetica':
        n_node = 43098
    elif opt.dataset == 'yoochoose1_64' or opt.dataset == 'yoochoose1_4':
        n_node = 37484
    else:
        n_node = 310

    # 获取训练模型
    model = trans_to_cuda(SessionGraph(opt, n_node))

    # 开始时间
    start = time.time()
    best_result = [0, 0]
    best_epoch = [0, 0]
    bad_counter = 0
    for epoch in range(opt.epoch):
        print('-------------------------------------------------------')
        print('epoch: ', epoch)
        # 评价指标：HR 命中率和 MRR 平均倒数排名
        hit, mrr = train_test(model, train_data, test_data)
        flag = 0
        # 更新最好的结果，如果此次训练结果好于之前结果，将结果储存，置flag为1（每次训练开始置flag为0）
        if hit >= best_result[0]:
            best_result[0] = hit
            best_epoch[0] = epoch
            flag = 1
        if mrr >= best_result[1]:
            best_result[1] = mrr
            best_epoch[1] = epoch
            flag = 1
        # 输出最好的结果
        print('Best Result:')
        print('\tRecall@20:\t%.4f\tMMR@20:\t%.4f\tEpoch:\t%d,\t%d' % (
            best_result[0], best_result[1], best_epoch[0], best_epoch[1]))
        bad_counter += 1 - flag
        # 如果多次训练效果没有提升则结束训练，默认超过10次出现训练效果没有提升的情况则结束训练
        if bad_counter >= opt.patience:
            break
    print('-------------------------------------------------------')
    end = time.time()
    print("Run time: %f s" % (end - start))


if __name__ == '__main__':
    main()
